from __future__ import division
import os
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import xml.dom.minidom
import numpy as np

ImgPath = '/fs/scratch/Szh6_XC-DA_edy_cv_gpu_users/waymo/waymococo_full/train2020/'
AnnoPath = '/fs/scratch/Szh6_XC-DA_edy_cv_gpu_users/waymo/waymococo_full/annotations/train2020_xml/'
ProcessedPath = 'waymo_crop_train'
if not os.path.exists(ProcessedPath):
    os.mkdir(ProcessedPath)

imagelist = os.listdir(ImgPath)

for image in imagelist:
    image_pre, ext = os.path.splitext(image)
    imgfile = ImgPath + image
    print(imgfile)
    if not os.path.exists(AnnoPath + image_pre + '.xml' ):
        continue
    xmlfile = AnnoPath + image_pre + '.xml'
    DomTree = xml.dom.minidom.parse(xmlfile)
    annotation = DomTree.documentElement
    filenamelist = annotation.getElementsByTagName('filename')
    filename = filenamelist[0].childNodes[0].data
    img_width = float(annotation.getElementsByTagName("width")[0].childNodes[0].data)
    img_height = float(annotation.getElementsByTagName("height")[0].childNodes[0].data)
    # print(img_width, img_height)
    objectlist = annotation.getElementsByTagName('object')
    i = 1
    for objects in objectlist:
        namelist = objects.getElementsByTagName('name')
        objectname = namelist[0].childNodes[0].data
        savepath = objectname        
        if not os.path.exists(ProcessedPath + '/' + savepath):
            os.makedirs(ProcessedPath + '/' + savepath)
        bndbox = objects.getElementsByTagName('bndbox')
        cropboxes = []
        for box in bndbox:
            x1_list = box.getElementsByTagName('xmin')
            x1 = float(x1_list[0].childNodes[0].data)
            y1_list = box.getElementsByTagName('ymin')
            y1 = float(y1_list[0].childNodes[0].data)
            x2_list = box.getElementsByTagName('xmax')
            x2 = float(x2_list[0].childNodes[0].data)
            y2_list = box.getElementsByTagName('ymax')
            y2 = float(y2_list[0].childNodes[0].data)
            if x1 <= 0:
                x1 = 1
            if x1 >= img_width:
                x1 = img_width -1
            if y1 <= 0:
                y1 = 1
            if y1 >= img_height:
                y1 = img_height -1
            if x2 <= 0:
                x2 = 1
            if x2 >= img_width:
                x2 = img_width -1
            if y2 <= 0:
                y2 = 1
            if y2 >= img_height:
                y2 = img_height -1
            if x1 >= x2:
                x2 = x1 +1
            if y1 >= y2:
                y2 = y1 +1
            w = x2 - x1
            h = y2 - y1
            obj = np.array([x1,y1,x2,y2])
            shift = np.array([[1,1,1,1]])
            XYmatrix = np.tile(obj,(1,1))
            cropboxes = XYmatrix * shift
            img = Image.open(imgfile)
            for cropbox in cropboxes:
                cropedimg = img.crop(cropbox)
                cropedimg.save((ProcessedPath + '/' + savepath + '/' + image_pre + '_' + str(i) + '.jpg'), quality=95, subsampling=0)
                i += 1